#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time : 2024/11/13 12:49
# @Author : XXX
# @Site : 
# @File : train_best_model.py
# @Software: PyCharm

import sys
import os
import time
from collections import deque
from datetime import datetime

import torch
import torch.nn as nn
import torch.optim as optim
# import torch.nn.functional as F
import random
import math
import numpy as np

import pandas as pd
import json

import matplotlib.pyplot as plt

#from model.model_test import TestingInfo, TestingInfoCdu, TestingInfoFsu
from util.loguru_util import myLogger


from train_all_offline import  process_cdudata_neural, process_longrundata_neural, process_fsudata_neural
from train_all_offline import train_all_offline

from all_config import need_all_key


#列表第一个元素为longrun文件，第二个为cdu，第三个为fsu
###############################训练文件参数
all_file_list = [
    r'J:\1液冷待上传代码文件文件夹\新建文件夹\longrun1128.xlsx',
    r'J:\1液冷待上传代码文件文件夹\新建文件夹\cdu1128.xlsx',
    r'J:\1液冷待上传代码文件文件夹\新建文件夹\fsu1128.xlsx'
]

# try:
#     os.remove('memory.pkl')
# except:
#     myLogger.info(f"没有这个文件")
###############################测试文件参数
test_all_file_list = [
    r'F:\OneDrive - stu.xpu.edu.cn\烽火实习\yeleng_data\yeleng\测试\longrun_test_400v1125.xlsx',
    r'F:\OneDrive - stu.xpu.edu.cn\烽火实习\yeleng_data\yeleng\测试\cdu_test_400v1125.xlsx',
    r'F:\OneDrive - stu.xpu.edu.cn\烽火实习\yeleng_data\yeleng\测试\fsu_test_400v1125.xlsx'
]





# 评估代码中是否需要进行画图
is_draw = True
### 组装日期
current_datetime = datetime.now()
current_date = current_datetime.date().strftime("v%m_%d")
# 文件路径
file_name = 'model'


#超参数组合列表
epoch = 50
error_avg_set = 0.5

# batchSize_list = [8, 16]
# lr_list = [0.05, 0.005]
# dropout_rate_list = [0.08, 0.05]

batchSize_list = [16, 32]
lr_list = [0.001, 0.0005]


search_best_para = [(batchSize, lr)
                        # for dryValve in [0, 100]
                        for batchSize in batchSize_list
                        for lr in lr_list]

    ####################### 训练部分代码
for param in search_best_para:
    myLogger.info(f"当前参数组合为:{param}")
    time.sleep(1)
    model = train_all_offline(all_file_list, epoch, param[0], param[1], need_all_key=need_all_key)
    # 设置模型为评估模式
    model.eval()


    ################# 评估部分代码

    file_path_longrun = test_all_file_list[0]
    df_longrun = pd.read_excel(file_path_longrun)
    myLogger.info(f'read longrun data completed!')
    file_path_cdu = test_all_file_list[1]
    df_cdu = pd.read_excel(file_path_cdu)
    myLogger.info(f'read cdu data completed!')
    file_path_fsu = test_all_file_list[2]
    df_fsu = pd.read_excel(file_path_fsu)
    myLogger.info(f'read fsu data completed!')

    new_all_sorted_dict = {}


    device = torch.device(
        "cuda" if torch.cuda.is_available() else "cpu")
    prediction_1 = []
    real = []
    line_read_predit = df_cdu.shape[0]

    deque_save_pre = deque([50, 50, 50, 50], maxlen=4)

    for i in range(1, line_read_predit):
        row_longrun = df_longrun.iloc[i]
        json_str_longrun = row_longrun.to_json(date_format='iso', force_ascii=False)
        longrun_data = json.loads(json_str_longrun)
        # print("json_str_longrun:",longrun_data)
        row_cdu = df_cdu.iloc[i]
        json_str_cdu = row_cdu.to_json(date_format='iso', force_ascii=False)
        cdu_data = json.loads(json_str_cdu)
        # print("json_str_cdu:", cdu_data)
        row_fsu = df_fsu.iloc[i]
        json_str_fsu = row_fsu.to_json(date_format='iso', force_ascii=False)
        fsu_data = json.loads(json_str_fsu)
        #print("json_str_sfu:", fsu_data)

        # 这里对重要状态进行筛选,带有neural的是用于进行智能的状态值的输入（离线训练）
        # longrun_sorted_dict = process_longrundata(longrun_data)
        longrun_sorted_dict_neural = process_longrundata_neural(longrun_data)
        # myLogger.info(flen(longrun_sorted_dict): {len(longrun_sorted_dict)}")
        # myLogger.info(f'longrun_sorted_dict: {longrun_sorted_dict}')
        # cdu_sorted_dict = process_cdudata(cdu_data)
        cdu_sorted_dict_neural = process_cdudata_neural(cdu_data)

        # fsu_sorted_dict = process_fsudata(fsu_data)
        fsu_sorted_dict_neural = process_fsudata_neural(fsu_data)

        # all_sorted_dict = {**longrun_sorted_dict_neural, **cdu_sorted_dict_neural, **fsu_sorted_dict_neural}
        all_sorted_dict = {**longrun_sorted_dict_neural, **cdu_sorted_dict_neural, **fsu_sorted_dict_neural}

        # myLogger.info(f"len(longrun_pddata): {len(longrun_pddata)}")
        # longrun_pddata = pd.Series(longrun_sorted_dict)
        # cdu_pddata = pd.Series(cdu_sorted_dict)
        # fsu_pddata = pd.Series(fsu_sorted_dict)

        # longrun_pddata_neural = pd.Series(longrun_sorted_dict_neural)
        # cdu_pddata_neural = pd.Series(cdu_sorted_dict_neural)
        # fsu_pddata_neural = pd.Series(fsu_sorted_dict_neural)


        all_pddata = pd.Series(all_sorted_dict)


        ###########################################################################################################
        need_all_key = need_all_key

        # 数据预处理部分
        for key in need_all_key:
            if key == 'dryFan':

                if all_sorted_dict[key] > 65530:
                    all_sorted_dict[key] = int(sum(deque_save_pre) / len(deque_save_pre))
                deque_save_pre.append(all_sorted_dict[key])
            new_all_sorted_dict[key] = all_sorted_dict[key]

        # 组装标签
        label_dict = {'cduTso': all_pddata['cduTso']}
        new_all_sorted_dict.pop('cduTso')
        # 输入特征
        all_pddata = pd.Series(new_all_sorted_dict)
        label_pddata = pd.Series(label_dict)
        all_pddata = list(all_pddata)
        label_pddata = list(label_pddata)

        # 转换为张量
        xx = torch.tensor(all_pddata, dtype=torch.float, requires_grad=True).to(device)
        yy = torch.tensor(label_pddata, dtype=torch.float, requires_grad=True).to(device)

        prediction = model(xx)  # 前向传播，调用刚写好的序列函数
        prediction_1.append(prediction.item())
        real.append(yy.item())


    # 评测模型性能的相关参数
    if is_draw:
        plt.plot(prediction_1, label='prediction', color='g')  # 第三条线
        plt.plot(real, label='real', color='y')  # 第三条线
        plt.legend()
        plt.show()

    sub_real_pre = [abs(a - b) for a, b in zip(prediction_1, real)]
    Avg_error = sum(sub_real_pre)/len(sub_real_pre)
    myLogger.info(f"平均预测温度误差为{Avg_error}")
    myLogger.info(f"预测温度误差为{sub_real_pre}")
    is_need_hand_save = int(input('手动存储则输入1，否则输入0'))
    if (Avg_error < error_avg_set) or is_need_hand_save:
        myLogger.error(f"这是一个好模型")
        save_model_name = "Model_pre{}.pth".format(current_date)
        save_all_name = file_name + '\\' + save_model_name
        torch.save(model.state_dict(), save_all_name)
        myLogger.error(f"好模型已经保存成功")
        myLogger.info(f"参数组合为：{param}")
        break
    else:
        myLogger.info(f"该模型不满足要求，重新开始训练")



